from llmtcg.agents.testers.os.linux import LinuxSyscallTesterAgent
from llmtcg.agents.coders.c import CCoderAgent, CodingFailed
from llmtcg.chat.models.claude import ClaudeChatModel
from llmtcg.chat.models.local import LocalOpenaiChatModel
from llmtcg.chat.models.deepseek import DeepSeekChatModel
from llmtcg.chat.models.groq import GroqChatModel
from llmtcg.chat.models.mistral import MistralChatModel
from llmtcg.chat.models.openai import OpenaiChatModel
from llmtcg.chat.models.yi import YiChatModel
from llmtcg.utils.profiler import Profiler

from loguru import logger

from subprocess import run, PIPE
from argparse import ArgumentParser, Namespace
from typing import Optional, Tuple, List
from os import path, makedirs
from shutil import rmtree
from sys import exit

import multiprocessing
import atexit
import json


# Facts about the models
# - For code generation, even claude3-opus generates problamatic from time to time, while it
#   is really slow (and expensive!) compared to haiku


profiler = Profiler()

# coder_model = LocalOpenaiChatModel("http://127.0.0.1:8000/v1", "mistral")
# coder_model = YiChatModel(model=YiChatModel.Model.YI_SPARK)
coder_model = OpenaiChatModel(model=OpenaiChatModel.Model.GPT_35_TURBO)
coder_model.profiler = profiler

tester = LinuxSyscallTesterAgent(coder_model)
coder = CCoderAgent(coder_model)
coder.compiler_args.append("-Werror")

_shared_counter = multiprocessing.Manager().Queue()
_shared_total = 0


def parse_args() -> Namespace:
    parser = ArgumentParser(
        "python3 llmtcg.apps.linux_syscall_tcg",
        description="generating test cases for linux kernel syscalls"
    )
    parser.add_argument(
        "--syscalls", type=str,
        help="list of syscall names, separated by ',', e.g. 'open,read'. " +
             "By default all syscalls are generated"
    )
    parser.add_argument(
        "--resume", action="store_true",
        help="generate test code for previously failed scenes"
    )
    parser.add_argument(
        "--use-manpage", action="store_true",
        help="use linux manpage to grab syscall documents"
    )
    parser.add_argument(
        "--use-coding-history", action="store_true",
        help="use coding history to improve code generation performance"
    )
    parser.add_argument(
        "--jobs", metavar='N', type=int, default=1,
        help="using multiple works to generate test case in parallel"
    )
    parser.add_argument(
        "-o", "--output", type=str, required=True,
        help="the directory to dump the generated test cases"
    )
    return parser.parse_args()


def initialize(args: Namespace) -> None:
    coder.with_history = args.use_coding_history
    logger.info(f"use coding history: {coder.with_history}")


def get_manpage(syscall: str) -> Optional[str]:
    proc = run(f"man {syscall}", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    if proc.returncode != 0:
        logger.warning(f"failed to get manpage of syscall {syscall}")
        return None

    content = proc.stdout.decode()
    return content


def generate_test_code(scene: str, target_c_path: str) -> None:
    index = _shared_counter.qsize()
    try:
        logger.debug(f"generating test code for {scene}")
        program = coder.get_program(f"""
        Generate a c testing program that implements the following test scene: {scene}.
        The testing program should:

        1. directly use the assert() function to give the assertions and report errors,
        2. never try to invoke shell commands through `system` or simular functions,
        3. try to invoke syscalls through their corresponding libc functions, or the
           'syscall' function from the header file 'sys/syscall.h',
        3. do noy rely on any user input, in fact, you should assume that the user give
           no input from stdin,
        """)
        logger.info(f"({index}/{_shared_total}) succeeded to generate {target_c_path}")

        with open(target_c_path, "w+") as fprog:
            fprog.write(program)

        with open(target_c_path + ".log", "w+") as flog:
            summary_json = coder.last_context.summary_json()
            flog.write(json.dumps(summary_json))

        _shared_counter.put((target_c_path, True))

    except CodingFailed as ex:
        logger.error(f"({index}/{_shared_total}) failed to generate {target_c_path}: {ex}")
        with open(f"{target_c_path}.err", "w+") as ferr:
            ferr.write(f"error: {ex}")
            ferr.write(coder.last_context.summary())

        _shared_counter.put((target_c_path, False))


def generate_syscall(syscall: str, use_manpage: bool, target_dir: str) -> None:
    if path.exists(target_dir):
        logger.warning(f"will remove {target_dir} as it already exists")
        rmtree(target_dir)

    makedirs(target_dir, exist_ok=True)

    manpage = get_manpage(syscall)
    scene_prompt = ""
    if manpage is None and use_manpage:
        logger.error(f"cannot find manpage for {syscall}, it could not be a valid syscall")
        return
    elif manpage is not None and use_manpage:
        scene_prompt += f"""
        Generate a full list of test scenes of the linux syscall '{syscall}'
        based on the manpage shown below.

        ```
        {manpage}
        ```
        """
    else:
        scene_prompt += f"Generate a full list of test scenes of the linux syscall '{syscall}'."

    scene_prompt += """
    You need to create as many test scenes as possible, and work really hard to explore every
    potentially different behavior of this syscall.
    """

    scenes = tester.generate_test_scene(scene_prompt)
    logger.info(f"{len(scenes)} test scenes generated for syscall {syscall}")

    for i, scene in enumerate(scenes):
        # dumping the test scene
        with open(path.join(target_dir, f"test_case_{i}.txt"), "w+") as fdesc:
            fdesc.write(scene)

        c_program_path = path.join(target_dir, f"test_case_{i}.c")
        generate_test_code(scene, c_program_path)


def resume_generate(args: Namespace) -> None:
    global _shared_total

    target_dir = args.output
    nworkers = args.jobs
    jobs: List[Tuple[str, str]] = []

    proc = run(f"find {target_dir} | grep .txt$", shell=True, stdout=PIPE, stderr=PIPE)
    scene_paths = proc.stdout.decode().splitlines()
    logger.info(f"find {len(scene_paths)} test scenes in total")
    for scene_path in scene_paths:
        case_path = scene_path[:-4] + ".c"
        if path.isfile(case_path):
            logger.debug(f"test case of {scene_path} has already been generated")
        elif path.isfile(f"{case_path}.err"):
            logger.debug(f"ignore {scene_path} as it has encountered an error before")
        elif path.isfile(f"{scene_path}.ignore"):
            logger.debug(f"ignore {scene_path} as it is manually ignored")
        else:
            with open(scene_path, "r") as fscene:
                scene = fscene.read()

            jobs.append((scene, case_path))

    _shared_total = len(jobs)

    if nworkers == 1:
        for i, (scene, case_path) in enumerate(jobs):
            generate_test_code(scene, case_path)
    else:
        logger.info(f"using {nworkers} workers to generate test cases")
        pool = multiprocessing.Pool(nworkers)
        pool.starmap(generate_test_code, jobs)


def generate_syscalls(args: Namespace) -> None:
    if args.syscalls is None:
        syscall_names = tester.get_all_supported_syscalls()
    else:
        syscall_names = sorted(s.strip() for s in args.syscalls.split(","))

    output_dir = args.output
    logger.info(f"generating test cases for {syscall_names}")
    logger.info(f"there are {len(syscall_names)} syscalls in total")

    for syscall in syscall_names:
        syscall_target_dir = path.join(output_dir, syscall)
        generate_syscall(syscall, args.use_manpage, syscall_target_dir)


def print_stats() -> None:
    profiler.print_stats()


if __name__ == "__main__":
    args = parse_args()
    if not coder_model.validate():
        logger.critical(f"{coder_model.pretty_name()} is not properly configured")
        exit(1)

    atexit.register(print_stats)
    initialize(args)

    if args.resume:
        resume_generate(args)
    else:
        generate_syscalls(args)
